import torch
import logging
import os


def denorm(x: torch.tensor, mean, std):
    mean = torch.tensor(mean).view(1, 3, 1, 1).to(x.device)
    std = torch.tensor(std).view(1, 3, 1, 1).to(x.device)
    out = x * std + mean
    return out.clamp_(0, 1)


class Logger:
    def __init__(self, log_file_path):
        # 创建日志目录（如果不存在）
        os.makedirs(os.path.dirname(log_file_path), exist_ok=True)

        # 设置日志格式
        log_format = '%(asctime)s - %(levelname)s - %(message)s'

        # 创建一个logger
        self.logger = logging.getLogger('MyLogger')
        self.logger.setLevel(logging.DEBUG)

        # 创建一个控制台处理器
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.DEBUG)
        console_handler.setFormatter(logging.Formatter(log_format))

        # 创建一个文件处理器
        file_handler = logging.FileHandler(log_file_path)
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(logging.Formatter(log_format))

        # 将处理器添加到logger
        self.logger.addHandler(console_handler)
        self.logger.addHandler(file_handler)

    def log(self, message):
        self.logger.debug(message)


class EarlyStop:
    def __init__(self, patience=20, verbose=False):
        """
        :param patience: 多少个epoch内没有改进后停止训练
        :param verbose: 是否打印信息
        """
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False

    def __call__(self, metric):
        # 如果没有最好的分数，初始化
        if self.best_score is None:
            self.best_score = metric
        # 如果metric提高，更新最好的分数
        elif metric < self.best_score:
            self.best_score = metric
            self.counter = 0  # 重置计数
        else:
            self.counter += 1  # 计数器增加
            if self.counter >= self.patience:
                self.early_stop = True
                if self.verbose:
                    print(f"Early stopping triggered after {self.counter} epochs")